# -*- coding: utf-8 -*-
"""
@Time    : 2024/7/8 19:35 
@Author  : ZhangShenao 
@File    : bow_search_engine.py 
@Desc    : 基于Bag of Words模型的搜索引擎实现
"""
import re

from search_engine_base import SearchEngineBase


class BOWSearchEngine(SearchEngineBase):
    """基于Bag of Words模型的搜索引擎实现"""

    def __init__(self):
        """构造函数"""

        # 调用父类的构造函数
        super().__init__()

        # 初始化索引表
        self.__index_table = {}

    def process_corpus(self, corpus_id: str, text: str):
        """
        处理语料
        :param corpus_id: 语料ID
        :param text: 语料的文本内容
        """
        # 分词处理
        words = self.split_word(text)

        # 将分词结果存入索引表
        self.__index_table[corpus_id] = words

    def search(self, query: str) -> list[str]:
        """
        搜索功能实现
        :param query: 关键词
        :return: 关键词所在的语料ID列表
        """

        # 对query分词处理
        query_words = self.split_word(query)

        # 遍历索引表,判断是否query的每个分词都在语料文件中出现
        result = []
        for corpus_id, text in self.__index_table.items():
            if self.match_query_word(query_words, text):
                result.append(corpus_id)

        # 返回结果
        return result

    @staticmethod
    def split_word(text: str) -> set[str]:
        """
        对文本分词
        :param text: 元素文本
        :return: 分词后的单词集合
        """

        # 使用正则表达式,去除标点符号和换行符
        text = re.sub(r'[^\w ]', ' ', text)
        # 统一转换成小写
        text = text.lower()
        # 按照空格进行切分
        words = text.split(' ')
        # 去除空白单词
        words = filter(None, words)
        # 将单词列表转换成集合返回
        return set(words)

    @staticmethod
    def match_query_word(query_words: set[str], text: str) -> bool:
        """判断query的每个分词是否都在text中出现"""
        for query_word in query_words:
            if query_word not in text:
                return False

        return True
